"""
@Time: 2021/3/20 下午 10:26
@Author: jinzhuan
@File: argument_toolkit.py
@Desc: 
"""
import torch
from cognlp import *
import json

torch.cuda.set_device(4)
tokenize_toolkit = TokenizeToolkit()
words = tokenize_toolkit.run(
    'The true voodoo-worshipper attempts nothing of importance without certain sacrifices which are intended to propitiate his unclean gods.')

argument_toolkit = ArgumentToolkit(
    bert_model='bert-base-cased',
    model_path='../../../cognlp/data/fn/argument/toolkit/model.pt',
    device=torch.device("cuda:4"),
    device_ids=[4],
    max_seq_length=256,
    trigger_vocabulary_path='../../../cognlp/data/fn/argument/toolkit/frame_vocabulary.txt',
    argument_vocabulary_path='../../../cognlp/data/fn/argument/toolkit/argument_vocabulary.txt')
frames = [{'word': 'attempts', 'position': 3, 'frame': 'Attempt'}, {'word': 'importance', 'position': 6, 'frame': 'Importance'}, {'word': 'sacrifices', 'position': 9, 'frame': 'Rite'}, {'word': 'intended', 'position': 12, 'frame': 'Purpose'}]

argument_result = argument_toolkit.run(words, frames)
print(json.dumps(argument_result))
